package com.nbsaas.boot.sql;

import net.sf.jsqlparser.JSQLParserException;
import net.sf.jsqlparser.expression.Alias;
import net.sf.jsqlparser.expression.Expression;
import net.sf.jsqlparser.expression.LongValue;
import net.sf.jsqlparser.expression.operators.relational.ExpressionList;
import net.sf.jsqlparser.parser.CCJSqlParserUtil;
import net.sf.jsqlparser.schema.Column;
import net.sf.jsqlparser.statement.Statement;
import net.sf.jsqlparser.statement.select.*;

import java.util.ArrayList;
import java.util.Arrays;
import java.util.List;


/**
 * SQL查询转换为子查询并进行分组的工具类
 */
public class JsqlParserSubQuery {

    /**
     * 将SQL查询包装为子查询并添加分组
     *
     * @param originalSql 原始SQL查询
     * @param groupByColumns 分组字段列表
     * @param subQueryAlias 子查询别名，默认为 "sub"
     * @return 包装后的SQL字符串
     * @throws JSQLParserException 解析异常
     */
    public static String wrapAsSubQueryWithGroupBy(String originalSql, List<String> groupByColumns, String subQueryAlias) throws JSQLParserException {
        if (subQueryAlias == null || subQueryAlias.trim().isEmpty()) {
            subQueryAlias = "sub";
        }

        // 解析原始SQL
        Statement statement = CCJSqlParserUtil.parse(originalSql);
        if (!(statement instanceof Select)) {
            throw new IllegalArgumentException("只支持SELECT语句");
        }

        Select originalSelect = (Select) statement;

        // 创建新的外层查询
        PlainSelect outerSelect = new PlainSelect();

        // 设置子查询作为FROM表
        ParenthesedSelect subSelect = new ParenthesedSelect();
        subSelect.setSelect(originalSelect);
        subSelect.setAlias(new Alias(subQueryAlias));
        outerSelect.setFromItem(subSelect);

        // 添加SELECT字段 (从子查询别名中选择)
        List<SelectItem<?>> selectItems = new ArrayList<>();
        for (String groupColumn : groupByColumns) {
            Column column = new Column(subQueryAlias + "." + groupColumn);
            SelectItem<?> selectItem = new SelectItem<>(column);
            selectItems.add(selectItem);
        }

        // 添加聚合函数示例 (可根据需要修改)
        Column countColumn = new Column("*");
        ExpressionList<Expression> expressionList = new ExpressionList<>(countColumn);
        net.sf.jsqlparser.expression.Function countFunction = new net.sf.jsqlparser.expression.Function();
        countFunction.withName("count");
        countFunction.withParameters(expressionList);

        SelectItem<?> countItem = new SelectItem<>(countFunction, new Alias("total_count"));
        selectItems.add(countItem);

        outerSelect.setSelectItems(selectItems);

        // 添加GROUP BY子句
        if (!groupByColumns.isEmpty()) {
            GroupByElement groupBy = new GroupByElement();
            ExpressionList<Expression> groupByExpressions = new ExpressionList<>();
            for (String groupColumn : groupByColumns) {
                groupByExpressions.add(new Column(subQueryAlias + "." + groupColumn));
            }
            groupBy.setGroupByExpressions(groupByExpressions);
            outerSelect.setGroupByElement(groupBy);
        }

        // 返回外层查询的字符串表示
        return outerSelect.toString();
    }

    /**
     * 重载方法，使用默认别名 "sub"
     */
    public static String wrapAsSubQueryWithGroupBy(String originalSql, List<String> groupByColumns) throws JSQLParserException {
        return wrapAsSubQueryWithGroupBy(originalSql, groupByColumns, "sub");
    }

    /**
     * 重载方法，支持字符串数组形式的分组字段
     */
    public static String wrapAsSubQueryWithGroupBy(String originalSql, String[] groupByColumns, String subQueryAlias) throws JSQLParserException {
        return wrapAsSubQueryWithGroupBy(originalSql, Arrays.asList(groupByColumns), subQueryAlias);
    }

    /**
     * 重载方法，支持字符串数组且使用默认别名
     */
    public static String wrapAsSubQueryWithGroupBy(String originalSql, String[] groupByColumns) throws JSQLParserException {
        return wrapAsSubQueryWithGroupBy(originalSql, Arrays.asList(groupByColumns), "sub");
    }

    /**
     * 添加自定义聚合函数的版本
     *
     * @param originalSql 原始SQL查询
     * @param groupByColumns 分组字段
     * @param aggregateFunctions 聚合函数配置 (函数名 -> 字段名)
     * @param subQueryAlias 子查询别名
     * @return 包装后的SQL字符串
     */
    public static String wrapAsSubQueryWithCustomAggregation(String originalSql, List<String> groupByColumns,
                                                             List<QueryFunc> aggregateFunctions, String subQueryAlias) throws JSQLParserException {
        if (subQueryAlias == null || subQueryAlias.trim().isEmpty()) {
            subQueryAlias = "sub";
        }

        // 解析原始SQL
        Statement statement = CCJSqlParserUtil.parse(originalSql);
        if (!(statement instanceof Select)) {
            throw new IllegalArgumentException("只支持SELECT语句");
        }

        Select originalSelect = (Select) statement;

        // 创建新的外层查询
        PlainSelect outerSelect = new PlainSelect();

        // 设置子查询作为FROM表
        ParenthesedSelect subSelect = new ParenthesedSelect();
        subSelect.setSelect(originalSelect);
        subSelect.setAlias(new Alias(subQueryAlias));
        outerSelect.setFromItem(subSelect);

        // 添加SELECT字段
        List<SelectItem<?>> selectItems = new ArrayList<>();

        // 添加分组字段
        for (String groupColumn : groupByColumns) {
            Column column = new Column(subQueryAlias + "." + groupColumn);
            SelectItem<?> selectItem = new SelectItem<>(column);
            selectItems.add(selectItem);
        }

        // 添加自定义聚合函数
        for (QueryFunc aggFunc : aggregateFunctions) {
            Expression columnExpr;
            if ("*".equals(aggFunc.getField())) {
                columnExpr = new Column("*");
            } else {
                columnExpr = new Column(subQueryAlias + "." + aggFunc.getField());
            }

            ExpressionList<Expression> functionArgs = new ExpressionList<>(columnExpr);
            net.sf.jsqlparser.expression.Function function = new net.sf.jsqlparser.expression.Function();

            function.withName(aggFunc.getFunc());
            function.withParameters(functionArgs);

            Expression finalExpr = function;

            // 如果需要保留小数位数，用 ROUND 包裹
            if (("AVG".equalsIgnoreCase(aggFunc.getFunc())
                    || "SUM".equalsIgnoreCase(aggFunc.getFunc())
                    || "MAX".equalsIgnoreCase(aggFunc.getFunc())
                    || "MIN".equalsIgnoreCase(aggFunc.getFunc()))) {

                net.sf.jsqlparser.expression.Function roundFunc = new net.sf.jsqlparser.expression.Function();
                roundFunc.setName("ROUND");

                List<Expression> roundArgs = new ArrayList<>();
                roundArgs.add(function);
                if (aggFunc.getScale() == null) {
                    aggFunc.setScale(2);
                }

                roundArgs.add(new LongValue(aggFunc.getScale())); // 保留位数

                roundFunc.setParameters(new ExpressionList<>(roundArgs));

                finalExpr = roundFunc;
            }

            Alias alias = null;
            if (aggFunc.getAlias() != null && !aggFunc.getAlias().trim().isEmpty()) {
                alias = new Alias(aggFunc.getAlias());
            }
            SelectItem<?> selectItem = new SelectItem<>(finalExpr, alias);
            selectItems.add(selectItem);
        }

        outerSelect.setSelectItems(selectItems);

        // 添加GROUP BY子句
        if (!groupByColumns.isEmpty()) {
            GroupByElement groupBy = new GroupByElement();
            ExpressionList<Expression> groupByExpressions = new ExpressionList<>();
            for (String groupColumn : groupByColumns) {
                groupByExpressions.add(new Column(subQueryAlias + "." + groupColumn));
            }
            groupBy.setGroupByExpressions(groupByExpressions);
            outerSelect.setGroupByElement(groupBy);
        }

        // 返回外层查询的字符串表示
        return outerSelect.toString();
    }


    public static String summary(String originalSql, List<String> groupByColumns,
                                                             List<QueryFunc> aggregateFunctions) throws JSQLParserException {

        // 解析原始SQL
        Statement statement = CCJSqlParserUtil.parse(originalSql);
        if (!(statement instanceof Select)) {
            throw new IllegalArgumentException("只支持SELECT语句");
        }

        PlainSelect originalSelect = (PlainSelect) statement;
        //originalSelect.setAlias(new Alias("model"));

        // 添加SELECT字段
        List<SelectItem<?>> selectItems = new ArrayList<>();

        // 添加分组字段
        for (String groupColumn : groupByColumns) {
            Column column = new Column(groupColumn);
            SelectItem<?> selectItem = new SelectItem<>(column);
            selectItems.add(selectItem);
        }

        // 添加自定义聚合函数
        for (QueryFunc aggFunc : aggregateFunctions) {
            Expression columnExpr;
            if ("*".equals(aggFunc.getField())) {
                columnExpr = new Column("*");
            } else {
                columnExpr = new Column(aggFunc.getField());
            }

            ExpressionList<Expression> functionArgs = new ExpressionList<>(columnExpr);
            net.sf.jsqlparser.expression.Function function = new net.sf.jsqlparser.expression.Function();

            function.withName(aggFunc.getFunc());
            function.withParameters(functionArgs);

            Expression finalExpr = function;

            // 如果需要保留小数位数，用 ROUND 包裹
            if (("AVG".equalsIgnoreCase(aggFunc.getFunc())
                    || "SUM".equalsIgnoreCase(aggFunc.getFunc())
                    || "MAX".equalsIgnoreCase(aggFunc.getFunc())
                    || "MIN".equalsIgnoreCase(aggFunc.getFunc()))) {

                net.sf.jsqlparser.expression.Function roundFunc = new net.sf.jsqlparser.expression.Function();
                roundFunc.setName("ROUND");

                List<Expression> roundArgs = new ArrayList<>();
                roundArgs.add(function);
                if (aggFunc.getScale() == null) {
                    aggFunc.setScale(2);
                }

                roundArgs.add(new LongValue(aggFunc.getScale())); // 保留位数

                roundFunc.setParameters(new ExpressionList<>(roundArgs));

                finalExpr = roundFunc;
            }

            Alias alias = null;
            if (aggFunc.getAlias() != null && !aggFunc.getAlias().trim().isEmpty()) {
                alias = new Alias(aggFunc.getAlias());
            }
            SelectItem<?> selectItem = new SelectItem<>(finalExpr, alias);
            selectItems.add(selectItem);
        }

        originalSelect.setSelectItems(selectItems);

        // 添加GROUP BY子句
        if (!groupByColumns.isEmpty()) {
            GroupByElement groupBy = new GroupByElement();
            ExpressionList<Expression> groupByExpressions = new ExpressionList<>();
            for (String groupColumn : groupByColumns) {
                groupByExpressions.add(new Column( groupColumn));
            }
            groupBy.setGroupByExpressions(groupByExpressions);
            originalSelect.setGroupByElement(groupBy);
        }
        originalSelect.setOrderByElements(null);
        // 返回外层查询的字符串表示
        return originalSelect.toString();
    }
    /**
     * 检测SQL语句中是否包含CTE（Common Table Expression）
     *
     * @param sql SQL语句
     * @return 如果包含CTE返回true，否则返回false
     * @throws JSQLParserException 解析异常
     */
    public static boolean hasCTE(String sql) throws JSQLParserException {
        Statement statement = CCJSqlParserUtil.parse(sql);
        if (!(statement instanceof Select)) {
            return false;
        }

        Select select = (Select) statement;
        return select.getWithItemsList() != null && !select.getWithItemsList().isEmpty();
    }
    public static String summary(String originalSql, List<String> groupByColumns,
                                 List<QueryFunc> aggregateFunctions, String subQueryAlias) throws JSQLParserException {
        if (subQueryAlias == null || subQueryAlias.trim().isEmpty()) {
            subQueryAlias = "sub";
        }

        // 解析原始SQL
        Statement statement = CCJSqlParserUtil.parse(originalSql);
        if (!(statement instanceof Select)) {
            throw new IllegalArgumentException("只支持SELECT语句");
        }

        PlainSelect originalSelect = (PlainSelect) statement;
        originalSelect.setAlias(new Alias(subQueryAlias));

        // 添加SELECT字段
        List<SelectItem<?>> selectItems = new ArrayList<>();

        // 添加分组字段
        for (String groupColumn : groupByColumns) {
            Column column = new Column(subQueryAlias + "." + groupColumn);
            SelectItem<?> selectItem = new SelectItem<>(column);
            selectItems.add(selectItem);
        }

        // 添加自定义聚合函数
        for (QueryFunc aggFunc : aggregateFunctions) {
            Expression columnExpr;
            if ("*".equals(aggFunc.getField())) {
                columnExpr = new Column("*");
            } else {
                columnExpr = new Column(subQueryAlias + "." + aggFunc.getField());
            }

            ExpressionList<Expression> functionArgs = new ExpressionList<>(columnExpr);
            net.sf.jsqlparser.expression.Function function = new net.sf.jsqlparser.expression.Function();

            function.withName(aggFunc.getFunc());
            function.withParameters(functionArgs);

            Expression finalExpr = function;

            // 如果需要保留小数位数，用 ROUND 包裹
            if (("AVG".equalsIgnoreCase(aggFunc.getFunc())
                    || "SUM".equalsIgnoreCase(aggFunc.getFunc())
                    || "MAX".equalsIgnoreCase(aggFunc.getFunc())
                    || "MIN".equalsIgnoreCase(aggFunc.getFunc()))) {

                net.sf.jsqlparser.expression.Function roundFunc = new net.sf.jsqlparser.expression.Function();
                roundFunc.setName("ROUND");

                List<Expression> roundArgs = new ArrayList<>();
                roundArgs.add(function);
                if (aggFunc.getScale() == null) {
                    aggFunc.setScale(2);
                }

                roundArgs.add(new LongValue(aggFunc.getScale())); // 保留位数

                roundFunc.setParameters(new ExpressionList<>(roundArgs));

                finalExpr = roundFunc;
            }

            Alias alias = null;
            if (aggFunc.getAlias() != null && !aggFunc.getAlias().trim().isEmpty()) {
                alias = new Alias(aggFunc.getAlias());
            }
            SelectItem<?> selectItem = new SelectItem<>(finalExpr, alias);
            selectItems.add(selectItem);
        }

        originalSelect.setSelectItems(selectItems);

        // 添加GROUP BY子句
        if (!groupByColumns.isEmpty()) {
            GroupByElement groupBy = new GroupByElement();
            ExpressionList<Expression> groupByExpressions = new ExpressionList<>();
            for (String groupColumn : groupByColumns) {
                groupByExpressions.add(new Column(subQueryAlias + "." + groupColumn));
            }
            groupBy.setGroupByExpressions(groupByExpressions);
            originalSelect.setGroupByElement(groupBy);
        }

        // 返回外层查询的字符串表示
        return originalSelect.toString();
    }
    /**
     * 示例和测试方法
     */
    public static void main(String[] args) {
        try {

            String sql = """
                WITH sales_cte AS (
                  SELECT category, SUM(amount) AS total_sales
                  FROM orders
                  GROUP BY category
                ),
                sales_cte2 AS (
                  SELECT category, SUM(amount) AS total_sales
                  FROM orders
                  GROUP BY category
                )
                SELECT category, total_sales
                FROM sales_cte order by category desc
                """;



            // 示例1：基本的分组查询
            String originalSql = "SELECT id, name, age, department FROM employees WHERE age > 25";
            List<String> groupColumns = Arrays.asList("department");

            String result1 = JsqlParserSubQuery.wrapAsSubQueryWithGroupBy(originalSql, groupColumns);
            System.out.println("示例1 - 基本分组:");
            System.out.println(result1);
            System.out.println();

            // 示例2：多字段分组
            String multiGroupResult = wrapAsSubQueryWithGroupBy(originalSql, Arrays.asList("department", "age"));
            System.out.println("示例2 - 多字段分组:");
            System.out.println(multiGroupResult);
            System.out.println();

            QueryFunc func=  new QueryFunc();
            func.setScale(1);
            func.setAlias("je");
            func.setFunc("avg");
            func.setField("je");
            // 示例3：自定义聚合函数
            List<QueryFunc> aggregateFunctions = Arrays.asList(
                    func

            );

            String customAggResult = JsqlParserSubQuery.summary(
                    sql,
                    Arrays.asList(),
                    aggregateFunctions,"model"
            );
            System.out.println("示例3 - 自定义聚合函数:");
            System.out.println(customAggResult);

        } catch (JSQLParserException e) {
            System.err.println("SQL解析错误: " + e.getMessage());
            e.printStackTrace();
        }
    }
}
